#!/usr/bin/python3
import unittest

from network_flow import Graph, parse

class GraphTest(unittest.TestCase):

    def setUp(self):
        V, E, self.s, self.t = parse("testfiles/rail.txt")
        self.G = Graph(V, E)

    def tearDown(self):
        del self.G, self.s, self.t
    
    def test_max_flow(self):
        assert 163 == self.G.find_max_flow(self.s, self.t)
    
    def test_current_flow(self):
        assert self.G.current_flow(self.s) == self.G.current_flow(self.t)
        self.G.find_max_flow(self.s, self.t)
        assert self.G.current_flow(self.s) == self.G.current_flow(self.t)
        
    def test_min_cut(self):
        max_flow, min_cut = self.G.find_min_cut(self.s, self.t)
        
        assert len(min_cut) < len(self.G.E)
        min_cut_capacity = sum([edge.capacity for edge in min_cut])
        assert min_cut_capacity == max_flow, "%d != %s" %(min_cut_capacity, max_flow)

if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()